// MTPDevice.cpp : Defines the exported functions for the DLL application.
//

#include "stdafx.h"
#include "MTPDevice.h"
#include <comdef.h>
#include <mswmdm.h>
#include <scclient.h>
#include "mswmdm_i.c" 

#pragma comment(lib, "rsa32sdk2.lib")
#pragma comment(lib, "randlib.lib")
#pragma comment(lib, "bignumsdk.lib")
#pragma comment(lib, "Mssachlp.lib")

_COM_SMARTPTR_TYPEDEF(IComponentAuthenticate, __uuidof(IComponentAuthenticate));
_COM_SMARTPTR_TYPEDEF(IWMDeviceManager2, __uuidof(IWMDeviceManager2));
_COM_SMARTPTR_TYPEDEF(IWMDMEnumDevice, __uuidof(IWMDMEnumDevice));
_COM_SMARTPTR_TYPEDEF(IWMDeviceManager3, __uuidof(IWMDeviceManager3));

_COM_SMARTPTR_TYPEDEF(IWMDMDevice, __uuidof(IWMDMDevice));
_COM_SMARTPTR_TYPEDEF(IWMDMDevice2, __uuidof(IWMDMDevice2));
_COM_SMARTPTR_TYPEDEF(IWMDMDevice3, __uuidof(IWMDMDevice3));
_COM_SMARTPTR_TYPEDEF(IWMDMEnumStorage, __uuidof(IWMDMEnumStorage));
_COM_SMARTPTR_TYPEDEF(IWMDMStorage, __uuidof(IWMDMStorage));
_COM_SMARTPTR_TYPEDEF(IWMDMStorage2, __uuidof(IWMDMStorage2));
_COM_SMARTPTR_TYPEDEF(IWMDMStorageGlobals, __uuidof(IWMDMStorageGlobals));

CSecureChannelClient *pSacClient = NULL;

HRESULT WINAPI	MTPDeviceEnum(IWMDeviceManager2** ppWMDevMgr)
{
	HRESULT	hr = S_OK;
	IComponentAuthenticatePtr	spCompAuth;
	if(SUCCEEDED(hr = spCompAuth.CreateInstance(CLSID_MediaDevMgr)))
	{
		DWORD	dwCount = 0;
		BYTE abPVK[] = {0x00};    
		BYTE abCert[] = {0x00};  
		if(pSacClient == NULL)	pSacClient = new CSecureChannelClient();
		if(SUCCEEDED(hr = pSacClient->SetCertificate(SAC_CERT_V1, (BYTE*) abCert, sizeof(abCert), (BYTE*) abPVK, sizeof(abPVK))))
		{
			pSacClient->SetInterface(spCompAuth);
			if(SUCCEEDED(hr = pSacClient->Authenticate(SAC_PROTOCOL_V1)))
			{
				hr = spCompAuth->QueryInterface(IID_PPV_ARGS(ppWMDevMgr));
			}
		}
	}
	return	 hr;
}

HRESULT WINAPI MTPEncryptParam(BYTE* pBuffer, DWORD* bufLen, byte* pMac)
{
	HRESULT hr = S_OK;
	if(pSacClient == NULL) return E_FAIL;
	HMAC	hMac = NULL;
	if(SUCCEEDED(hr = pSacClient->MACInit(&hMac)))
	{
		hr = pSacClient->MACUpdate(hMac, pBuffer, *bufLen);
		hr = pSacClient->MACUpdate(hMac, (BYTE*)bufLen, sizeof(DWORD));
		hr = pSacClient->MACFinal(hMac, pMac);
		if(FAILED(hr)) return E_FAIL;

		if(FAILED(hr = pSacClient->EncryptParam(pBuffer, *bufLen))) return E_FAIL;

		return S_OK;
	}
	return E_FAIL;
}

HRESULT WINAPI MTPDecryptParam(BYTE* pBuffer, DWORD* bufLen, byte* pMac)
{
	HRESULT hr = S_OK;
	if(pSacClient == NULL) return E_FAIL;
	HMAC hMAC;
	BYTE pTestMac[WMDM_MAC_LENGTH];
	if(FAILED(hr = pSacClient->DecryptParam(pBuffer, *bufLen))) return E_FAIL;
	if(SUCCEEDED(hr = pSacClient->MACInit(&hMAC)))
	{
		hr = pSacClient->MACUpdate(hMAC, pBuffer, *bufLen);
		hr = pSacClient->MACUpdate(hMAC, (BYTE*)bufLen, sizeof(DWORD));
		hr = pSacClient->MACFinal(hMAC, pTestMac);
		if((memcmp(pMac, pTestMac, WMDM_MAC_LENGTH) != 0) || FAILED(hr)) return WMDM_E_MAC_CHECK_FAILED;
		return S_OK;
	}
	return E_FAIL;
}